//#include <iostream>
//#include "HIRCode.h"
//#include <string>
#include <cstring>
#include <cctype>
#include <filesystem>
//#include <exception>
#include "HExceptions.h"

// #include "antlr4-runtime.h"
#include "HVisitor.h"
#include "HWarnings.h"
#include "HRV64Generator.h"

using namespace antlr4;

extern HWarnings warnings;

int error_code = 0;

// optimization level 1, 2, 3: 'K', 'T', 'Z'
const char Optcode[] = {'0', OPT_L1, OPT_L2, OPT_L3};

// global flags of command line options
HOptions opts;

// strings of all command line arguments
std::vector<std::string> cmd_args;

// Log output
std::ofstream log_output;

static const char *get_ofile_name(const char *srcf_name, bool is_ir = false);
static int set_option(const char *opt);

int main(int argc, const char* argv[])
{
	std::ifstream fin;
	std::istream *pis;
	std::ofstream fout;	// output file
	const char *src_file = NULL;
	ANTLRInputStream input;

	opts.exe_file = argv[0];

	cmd_args.push_back(std::string(argv[0]));

	while (--argc > 0)
	{
		++argv;
		// Record all command line arguments
		cmd_args.push_back(std::string(*argv));

		if ((*argv)[0] == '-')
		{
			switch (set_option(*argv))
			{
			case 1:
				if (--argc > 0)
				{
					set_option(*++argv);
					cmd_args.push_back(std::string(*argv));
				}
				else
				{
					std::cout << HXD1_Info << Normal_error <<
						"missing filename after \'" << *argv << '\'' << std::endl;
					return 1;
				}
				break;
			case -1:
				std::cout << HXD1_Info << Normal_error <<
					"unrecognized command line option \'" << *argv << '\'' << std::endl;
				return 1;
				break;
			}
		}
		else
			src_file = *argv;
	}

	/*
	 * If debug mode is enabled, we allow input code from `stdin`.
	 * Otherwise, input file must be specified.
	 */
	if (src_file == NULL)
	{
		if (opts.debug_level != 0)
			pis = &std::cin;
		else
		{
			std::cerr << HXD1_Info << Fatal_error << "no input file" << std::endl
				<< "compilation terminated." << std::endl;
			return error_code = int(HECode::FTLERR);
		}	
	}
	else
	{
		fin.open(src_file);
		if (!fin.is_open() || std::filesystem::is_directory(src_file))
		{
			std::cerr << HXD1_Info << Fatal_error << "cannot open file " <<
				src_file << std::endl;
			return error_code = int(HECode::FTLERR);
		}
		pis = &fin;
	}

	input = ANTLRInputStream(*pis);

	CACTLexer lexer(&input);
	CommonTokenStream tokens(&lexer);
	CACTParser parser(&tokens);

	HAnalysis visitor;
	CACTParser::Comp_unitContext* comp_unit_parse_tree;

	// Parser, do!
	comp_unit_parse_tree = parser.comp_unit();

	try {
		hxd1_ir.add_builtin(functable, BUILTIN_NUM);
		// Visitor, do! (Visit the parse tree.)
		visitor.visit(comp_unit_parse_tree);
		if (error_code != 0)
		{	// If syntax error is found...
			assert(error_code == int(HECode::STXERR));
			throw SyntaxException();
		}
		// Do necessary processes and optimizations to HXD1 IR
		hxd1_ir.post_process();
		// Show warnings
		warnings.show();

		// Initialize log file
		if (opts.log_file != nullptr)
		{
			log_output.open(opts.log_file);
			if (!log_output.is_open())
			{
				std::cerr << HXD1_Info << Fatal_error << "cannot open file "
					<< opts.log_file << std::endl;
				// Ignore log file
				opts.log_file = nullptr;
			}
			else
			{
				log_output << "This is the log file dumped by HXD1D:";
				for (auto arg : cmd_args)
					log_output << ' ' << arg;
				log_output << std::endl << std::endl;
			}
		}

		// Dump HXD1 IR (only HXD1C works) or RV64 code (HXD1D also works)
		if (opts.emit_ircode != 0)
		{	// dump HXD1 IR code!!!
			if (opts.out_file == nullptr)
				opts.out_file = get_ofile_name(src_file, true);
			fout.open(opts.out_file);
			if (!fout.is_open())
			{
				std::cerr << HXD1_Info << Fatal_error << "cannot open file " <<
					opts.out_file << std::endl;
				return error_code = int(HECode::FTLERR);
			}
			hxd1_ir.dump(fout, src_file);
			if (opts.debug_level > 0)
				std::cout << HXD1_Info << "IR code dumped to file " <<
					opts.out_file << std::endl;
		}
		else
		{	// start RV64 code generation!!!
			if (opts.use_32_inst != 0)
				HRV64Generator::integer_inst_size = RV64Opsz::W;
			hxd1_rv64gen.init_from_IR(hxd1_ir);
			hxd1_rv64gen.generate();
			if (opts.out_file == nullptr)
				opts.out_file = get_ofile_name(src_file);
			fout.open(opts.out_file);
			if (!fout.is_open())
			{
				std::cerr << HXD1_Info << Fatal_error << "cannot open file " <<
					opts.out_file << std::endl;
				return error_code = int(HECode::FTLERR);
			}
			hxd1_rv64gen.dump(fout);
		}
	}
	catch (SyntaxException &ex)
	{
		ex.display();
		error_code = int(HECode::STXERR);
	}
	catch (SemanticException &ex)
	{
		ex.display();
		error_code = int(HECode::SMTERR);
	}
	catch (UnsupportedSituation &ex)
	{
		ex.display();
		error_code = int(HECode::UNSSIT);
	}
	catch (std::exception &ex) {
		std::cerr << ex.what() << std::endl;
		error_code = int(HECode::UKNERR);
		/*
		 * When unknown error happens, try to dump the RV64 code
		 * for debugging.
		 */
		if (opts.emit_ircode == 0)
		{
			if (opts.out_file == nullptr)
				opts.out_file = get_ofile_name(src_file);
			fout.open(opts.out_file);
			if (!fout.is_open())
				std::cerr << HXD1_Info << Fatal_error << "cannot open file " <<
					opts.out_file << std::endl;
			else
				hxd1_rv64gen.dump(fout);
		}
	}

	// Notice of log
	if (opts.log_file != nullptr && opts.debug_level > 0)
		std::cout << HXD1_Info << "Log of HXD1D dumped to file "
			<< opts.log_file << std::endl;
	// Close output files
	if (fout.is_open())
		fout.close();
	if (log_output.is_open())
		log_output.clear();
	// Report unknown error
	if (error_code == int(HECode::UKNERR))
		std::cerr << Normal_error << "unknown"
			<< std::endl;

	return error_code;
}

// generate output file name according to source file name (and path)
static const char *get_ofile_name(const char *srcf_name, bool is_ir)
{
	size_t last_slash, dot_find_start = 1ul, last_dot;
	std::string ofname;

	if (srcf_name == nullptr)
		ofname = "hxd1";
	else
	{
		ofname = srcf_name;

		last_slash = ofname.rfind('/');
		// if '/' is found in the source file path and name...
		if (last_slash != std::string::npos)
			// +2 to avoid file name starting with '.' such as ".../.file"
			dot_find_start = last_slash + 2u;
		
		last_dot = ofname.find('.', dot_find_start);
		if (last_dot != std::string::npos)
			// '.' is found
			ofname = ofname.substr(0ul, last_dot);
	}
	ofname += (is_ir ? IR_FILE_SUFFIX : ASM_FILE_SUFFIX);

	return strdup(ofname.c_str());
}

/*
 * Set option specified by string `opt`.
 * Return 0 on success or -1 if the option is invalid.
 * Besides, returning 1 means that the next option is needed.
 */
static int set_option(const char *opt)
{
	/*
	 * Flag representing that the last option is an option
	 * that needs a file following it.
	 * 0: None;
	 * 1: `-o`;
	 * 2: `-l`
	 */
	static int flag_output = 0;

	if (flag_output == 0)
	{
		if (opt[0] != '-')
			// option must start with '-'
			goto err_opt;
	
		switch (opt[1])
		{
		case 'O':
			switch (opt[2])
			{
			case '0':
				opts.opt_level = 0; break;
			case '1': case OPT_L1: case '\0':
				opts.opt_level = 1; break;
			case '2': case OPT_L2:
				opts.opt_level = 2; break;
			case '3': case OPT_L3:
				opts.opt_level = 3; break;
			default:
				goto err_opt;
				break;
			}
			if (opts.opt_level > 0)
			{
				opts.opt_str_red = opts.opt_dce = opts.opt_cse
					= opts.opt_merge_const_expr = opts.opt_cfg_cleanup = 1;
				opts.no_alloc_stack = 1;
			}
			if (opts.opt_level > 1)
				opts.opt_gra = 1;
			break;
		case 'd':
			opts.no_check_main = 1;
			if (isdigit(opt[2]))
				opts.debug_level = opt[2] - '0';
			else
				opts.debug_level = 1;
			if (opts.debug_level != 0)
				std::cout << HXD1_Info << "Debug mode is enabled with level " 
					<< int(opts.debug_level) << std::endl;
			break;
		case 'e':
			if (strcmp(opt, "-emit-hxd1") == 0 || strcmp(opt, "-emit-llvm") == 0
				|| strcmp(opt, "-emit-IR") == 0)
				opts.emit_ircode = 1;
			else
				// Simplify the argument
				opts.emit_ircode = 1;
				//goto err_opt;
			break;
		case 'S':
			// Dump RV64 asm code
			opts.emit_ircode = 0;
			break;
		case 'h': // show help info
		{
			if (opt[2] == 'f')
			{
				std::cout << "The following options are for features"
					" such as optimizations:" << std::endl;
				std::cout << "\t-fdump-builtin\t\t"
					"Dump builtin functions in IR code." << std::endl;
				std::cout << "\t-fno-check-main\t\t"
					"Don't check type of main()." << std::endl;
				std::cout << "\t-fuse-32-inst\t\t"
					"Use 32-bit instructions for basic calculations." << std::endl;
				std::cout << "\t-fno-alloc-stack\t"
					"Don't always allocate stack space for local variables." << std::endl;
				std::cout << "\t-fstrength-reduce\t"
					"Perform simple strength reduction." << std::endl;
				std::cout << "\t-fdce\t\t\t"
					"Perform local dead code elimination (DCE)." << std::endl;
				std::cout << "\t-fcse\t\t\t"
					"Perform local common subexpression elimination (CSE)." << std::endl;
				std::cout << "\t-fmerge-const-expr\t"
					"Attempt to calculate constant expressions while compiling."
					<< std::endl;
				std::cout << "\t-fcfg-cleanup\t\t"
					"Attempt to simplify continuous jumps and remove unreachable code."
					<< std::endl;
				std::cout << "\t-fgra\t\t\t"
					"Perform global register allocation using graph coloring."
					<< std::endl;
			}
			else
			{
				std::cout << HXD1_Info << "a simple CACT compiler built at "
					__DATE__ << ", " << __TIME__ << std::endl;
				std::cout << "Usage: " << opts.exe_file << " [options] <file>" << std::endl;
				std::cout << "Options:" << std::endl;
				std::cout << "\t-d<level>\t\tEnable debug mode." << std::endl;
				std::cout << "\t-h\t\t\tShow basic help info." << std::endl;
				std::cout << "\t-hf\t\t\tShow help info of options "
					"starting with -f." << std::endl;
				std::cout << "\t-l <file>\t\tWrite HXD1D's compiling log into <file>."
					<< std::endl;
				std::cout << "\t-O<level>\t\tEnable optimization where level can be "
					"0, " << OPT_L1 << ", " << OPT_L2 << ", " << OPT_L3 << " or 0, 1, 2, 3."
					<< std::endl;
				std::cout << "\t-o <file>\t\tPlace the output into <file>." << std::endl;
				std::cout << "\t-emit-IR\t\tDump HXD1 IR code to a file." << std::endl;
				std::cout << "\t-S\t\t\tDump RV64 assembly code. (default)" << std::endl;
				std::cout << "If <file> is omitted and -d is specified, "
					"stdin will be used to get input." << std::endl;
			}
			exit(0);
			return 0;
		}
			break;
		case 'f':
		{
			opt += 2l;
			int flag_set = 1;
			if (strncmp(opt, "no-", 3ul) == 0)
			{
				opt += 3l;
				flag_set = 0;
			}
			if (strcmp(opt, "dump-builtin") == 0)
				opts.dump_builtin = flag_set;
			else if (strcmp(opt, "check-main") == 0)
				opts.no_check_main = !flag_set;
			else if (strcmp(opt, "use-32-inst") == 0)
				opts.use_32_inst = flag_set;
			else if (strcmp(opt, "alloc-stack") == 0)
				opts.no_alloc_stack = !flag_set;
			else if (strcmp(opt, "strength-reduce") == 0)
				opts.opt_str_red = flag_set;
			else if (strcmp(opt, "dce") == 0)
				opts.opt_dce = flag_set;
			else if (strcmp(opt, "cse") == 0)
				opts.opt_cse = flag_set;
			else if (strcmp(opt, "merge-const-expr") == 0)
				opts.opt_merge_const_expr = flag_set;
			else if (strcmp(opt, "cfg-cleanup") == 0)
				opts.opt_cfg_cleanup = flag_set;
			else if (strcmp(opt, "gra") == 0)
				opts.opt_gra = flag_set;
			else
				goto err_opt;
		}
			break;
		case 'l': case 'o':
			flag_output = (opt[1] == 'l' ? 2 : 1);
			return 1;
			break;
		default:
			goto err_opt;
			break;
		}
	}
	else
	{
		if (flag_output == 1)
			opts.out_file = opt;
		else if (flag_output == 2)
			opts.log_file = opt;
		flag_output = 0;
	}
	return 0;
err_opt:
	return -1;
}
